package com.mega.fe.util.klass;

import com.mega.fe.util.MEGA;
import com.mega.fe.util.MyUnsafe;
import com.mega.fe.util.UnsafeUtil;
import com.mega.fe.asm.ClassWriter;
import sun.misc.Unsafe;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodType;
import java.lang.invoke.VarHandle;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.ProtectionDomain;
import java.util.HashMap;
import java.util.Map;

public class KlassCreator {
    public static HashMap<String, ClassInfo> CREATED_CLASSES = new HashMap<>();

    public static void main(String[] args) throws Throwable {
        long time0 = System.currentTimeMillis();
        //System.out.println(1);

        stringToClass("""
                public final class Main {
                    public static void main(String[] args) {
                        System.out.println("Hello World!");
                    }
                }
                """, "Main").getDeclaredMethod("main", String[].class).invoke(null, (Object) new String[0]);
        System.out.printf("time:%s", System.currentTimeMillis()-time0);

    }

    public static Class<?> stringToClass(String input, String className) throws Throwable {
        String tmpDir = File.createTempFile("tmp", null).getParentFile().getAbsolutePath();
        File f = new File(tmpDir + File.separatorChar + className + ".java");
        FileWriter fw = new FileWriter(f);
        BufferedWriter writer = new BufferedWriter(fw);
        writer.write(input);
        writer.close();
        fw.close();
        ProcessBuilder builder = new ProcessBuilder("javac", f.getAbsolutePath());
        builder.start().waitFor();
        File f1 = new File(f.getParentFile().getAbsolutePath() + File.separatorChar + f.getName().replace(".java", "") + ".class");
        byte[] buf = Files.readAllBytes(Path.of(f1.getAbsolutePath()));

        MethodHandle defineClass0 =
                MEGA.IMPL_LOOKUP()
                        .findVirtual(UnsafeUtil.INTERNAL_UNSAFE_CLASS(), "defineClass0"
                                , MethodType.methodType(Class.class, String.class, byte[].class, int.class, int.class, ClassLoader.class, ProtectionDomain.class))
                .bindTo(UnsafeUtil.INTERNAL_UNSAFE_INSTANCE());
        return (Class<?>) defineClass0.invoke(className, buf, 0, buf.length, Thread.currentThread().getContextClassLoader(), MEGA.class.getProtectionDomain());
    }


    public static Class<?> dynamicCreate(String class_name, String package_name, String JAVA_SOURCE_CODE) throws Throwable {
        JavaStringCompiler compiler = new JavaStringCompiler();
        Map<String, byte[]> results = compiler.compile( class_name+ ".java", ezCode(JAVA_SOURCE_CODE));
        System.out.println(results.keySet());
        CREATED_CLASSES.put(package_name+"."+class_name, new ClassInfo(results));
        Class<?> clazz = compiler.loadClass(package_name+"."+class_name, results);
        return clazz;
    }
    public static String ezCode(String... strings) {
        StringBuilder result = new StringBuilder();
        for (String s : strings)
            result.append(s);
        return result.toString();
    }
    public record ClassInfo(Map<String, byte[]> map) {}

    public static void setThisClass(ClassWriter cw, int version, String name) throws Throwable {
        VarHandle this_class = MEGA.IMPL_LOOKUP().findVarHandle(ClassWriter.class, "thisClass", int.class);
        Class<?> SymbolTableClass = Class.forName("com.mega.fe.asm.SymbolTable");
        VarHandle symbol_table = MEGA.IMPL_LOOKUP().findVarHandle(ClassWriter.class, "symbolTable", SymbolTableClass);
        MethodHandle setMajorVersionAndClassName = MEGA.IMPL_LOOKUP().findVirtual(SymbolTableClass, "setMajorVersionAndClassName", MethodType.methodType(int.class, int.class, String.class));
        this_class.set(cw, setMajorVersionAndClassName.bindTo(symbol_table.get(cw)).invoke(version, name));
    }
    public static void setSuperClass(ClassWriter cw, String name) throws Throwable {
        VarHandle superClass = MEGA.IMPL_LOOKUP().findVarHandle(ClassWriter.class, "superClass", int.class);
        Class<?> SymbolTableClass = Class.forName("com.mega.fe.asm.SymbolTable");
        Class<?> SymbolClass = Class.forName("com.mega.fe.asm.Symbol");
        VarHandle symbol_table = MEGA.IMPL_LOOKUP().findVarHandle(ClassWriter.class, "symbolTable", SymbolTableClass);
        VarHandle index = MEGA.IMPL_LOOKUP().findVarHandle(SymbolClass, "index", int.class);
        MethodHandle addConstantClass = MEGA.IMPL_LOOKUP().findVirtual(SymbolTableClass, "addConstantClass", MethodType.methodType(SymbolClass, String.class));
        superClass.set(cw, index.get(addConstantClass.bindTo(symbol_table.get(cw)).invoke(name)));
    }

    public static Object VM() throws Throwable {
        Class<?> vmClass = Class.forName("sun.jvm.hotspot.runtime.VM");
        MethodHandle methodHandle = MEGA.IMPL_LOOKUP().findStatic(vmClass, "getVM", MethodType.methodType(vmClass));
        return methodHandle.invoke();
    }

    public static void addClass(ClassLoader loader, Class<?> target) throws Throwable {
        MethodHandle mh = MEGA.IMPL_LOOKUP().findVirtual(ClassLoader.class, "addClass", MethodType.methodType(void.class, Class.class));
        mh.bindTo(loader).invoke(target);
    }
}
